from itertools import chain

import paddle.fluid.layers as L
from paddle.fluid.initializer import Normal
from paddle.fluid.dygraph import Layer, LayerList, Sequential, Linear, ParameterList

from helm.static.models.layers import BN, GlobalAvgPool, Conv2d
from helm.dynamic.models.nas.darts.operations import OPS, ReLUConvBN, FactorizedReduce


class MixedOp(Layer):

    def __init__(self, primitives, C, stride):
        super().__init__()
        self._ops = LayerList()
        for primitive in primitives:
            op = OPS[primitive](C, stride)
            self._ops.append(op)

    def forward(self, x, weights):
        ss = []
        for i, op in enumerate(self._ops):
            ss.append(weights[i] * op(x))
        return sum(ss)


class Cell(Layer):

    def __init__(self, primitives, steps, multiplier, C_prev_prev, C_prev, C, reduction, reduction_prev):
        super().__init__()
        self.reduction = reduction

        if reduction_prev:
            self.preprocess0 = FactorizedReduce(C_prev_prev, C)
        else:
            self.preprocess0 = ReLUConvBN(C_prev_prev, C, 1)
        self.preprocess1 = ReLUConvBN(C_prev, C, 1)
        self._steps = steps
        self._multiplier = multiplier

        self._ops = LayerList()
        for i in range(self._steps):
            ops = LayerList()
            for j in range(2 + i):
                stride = 2 if reduction and j < 2 else 1
                op = MixedOp(primitives, C, stride)
                ops.append(op)
            self._ops.append(ops)

    def forward(self, s0, s1, weights):
        s0 = self.preprocess0(s0)
        s1 = self.preprocess1(s1)
        states = [s0, s1]
        for i in range(self._steps):
            s = sum(self._ops[i][j](h, weights[i][j]) for j, h in enumerate(states))
            states.append(s)

        return L.concat(states[-self._multiplier:], axis=1)


class Network(Layer):

    def __init__(self, primitives, C, layers=8, steps=4, multiplier=4, stem_multiplier=3, num_classes=10):
        super().__init__()
        self._C = C
        self._num_classes = num_classes
        self._layers = layers
        self._steps = steps
        self._multiplier = multiplier
        self._primitives = primitives

        C_curr = stem_multiplier * C
        self.stem = Sequential(
            Conv2d(3, C_curr, kernel_size=3, bias=False),
            BN(C_curr),
        )

        C_prev_prev, C_prev, C_curr = C_curr, C_curr, C
        self.cells = LayerList()
        reduction_prev = False
        for i in range(layers):
            if i in [layers // 3, 2 * layers // 3]:
                C_curr *= 2
                reduction = True
            else:
                reduction = False
            cell = Cell(primitives, steps, multiplier, C_prev_prev, C_prev, C_curr, reduction, reduction_prev)
            self.cells.append(cell)
            reduction_prev = reduction
            C_prev_prev, C_prev = C_prev, multiplier * C_curr

        self.avg_pool = GlobalAvgPool()
        self.classifier = Linear(C_prev, num_classes)

        self._initialize_alphas()

    def forward(self, input):
        s0 = s1 = self.stem(input)
        for i, cell in enumerate(self.cells):
            weights = self.alphas_reduce if cell.reduction else self.alphas_normal
            weights = [L.softmax(w, axis=0) for w in weights]
            s0, s1 = s1, cell(s0, s1, weights)
        out = self.avg_pool(s1)
        logits = self.classifier(out)
        return logits

    def _initialize_alphas(self):
        num_ops = len(self._primitives)

        self.alphas_normal = ParameterList([
            L.create_parameter([i + 2, num_ops], 'float32', default_initializer=Normal(0, 1e-3))
            for i in range(self._steps)
        ])
        self.alphas_reduce = ParameterList([
            L.create_parameter([i + 2, num_ops], 'float32', default_initializer=Normal(0, 1e-3))
            for i in range(self._steps)
        ])

    def model_parameters(self):
        ids = set(id(p) for p in self.arch_parameters())
        for p in self.parameters():
            if id(p) not in ids:
                yield p

    def arch_parameters(self):
        return chain(self.alphas_normal.parameters(), self.alphas_reduce.parameters())

def test_net():

    net = Network(8, 8, 4, 4, 1)